import os
import time
import subprocess
import tempfile
import docker
from celery import shared_task
from submission.models.submission_model import SubmissionModel
from execution.models.test_case_model import TestCaseModel

@shared_task
def judge_submission(submission_id):
    """
    评测提交的代码
    """
    submission = SubmissionModel.objects.get(id=submission_id)
    submission.result = 2  # 设置为评测中
    submission.save()

    try:
        # 获取题目的测试用例
        test_cases = TestCaseModel.objects.filter(problem=submission.problem)
        submission.total_test_cases = test_cases.count()
        submission.save()

        # 准备 Docker 客户端
        try:
            client = docker.DockerClient(base_url="unix:///home/edy/.docker/desktop/docker.sock")
            # 验证连接
            client.ping()
        except docker.errors.DockerException as e:
            print(f"Docker connection error: {str(e)}")
            submission.result = 6  # 运行时错误
            submission.compile_output = f"Docker connection error: {str(e)}"
            submission.save()
            return

        # 根据语言选择镜像和运行命令
        config = get_language_config(submission.language)
        if not config:
            raise ValueError(f"Unsupported language: {submission.get_language_display()}")

        # 创建临时目录存放代码和测试用例
        with tempfile.TemporaryDirectory() as temp_dir:
            # 写入代码文件
            code_file = os.path.join(temp_dir, config['file_name'])
            with open(code_file, 'w') as f:
                f.write(submission.code)

            # 创建容器
            container = client.containers.create(
                image=config['image'],
                command='tail -f /dev/null',  # 保持容器运行
                volumes={
                    temp_dir: {'bind': '/app', 'mode': 'rw'}
                },
                working_dir='/app',
                mem_limit=f'{config["memory_limit"]}m',
                nano_cpus=int(1e9),  # 限制使用1个CPU核心
                network_disabled=True  # 禁用网络访问
            )

            try:
                container.start()

                # 编译代码（如果需要）
                if 'compile_cmd' in config:
                    compile_result = container.exec_run(
                        config['compile_cmd'],
                        workdir='/app'
                    )
                    if compile_result.exit_code != 0:
                        submission.result = 5  # 编译错误
                        submission.compile_output = compile_result.output.decode('utf-8')
                        submission.save()
                        return

                # 运行测试用例
                max_time = 0
                max_memory = 0
                passed_cases = 0

                for test_case in test_cases:
                    # 写入输入文件
                    input_file = os.path.join(temp_dir, 'input.txt')
                    with open(input_file, 'w') as f:
                        f.write(test_case.input_data)

                    # 运行代码
                    start_time = time.time()
                    run_result = container.exec_run(
                        config['run_cmd'],
                        workdir='/app',
                        stdin=True,
                        stderr=True
                    )
                    time_used = int((time.time() - start_time) * 1000)  # 转换为毫秒

                    # 检查运行结果
                    if run_result.exit_code != 0:
                        submission.result = 6  # 运行时错误
                        submission.compile_output = run_result.output.decode('utf-8')
                        break

                    # 检查时间限制
                    if time_used > config['time_limit']:
                        submission.result = 7  # 时间超限
                        break

                    # 获取内存使用情况
                    stats = container.stats(stream=False)
                    memory_used = stats['memory_stats']['usage'] / (1024 * 1024)  # 转换为MB

                    # 检查内存限制
                    if memory_used > config['memory_limit']:
                        submission.result = 8  # 内存超限
                        break

                    # 更新最大时间和内存使用
                    max_time = max(max_time, time_used)
                    max_memory = max(max_memory, memory_used)

                    # 比较输出
                    actual_output = run_result.output.decode('utf-8').strip()
                    expected_output = test_case.output_data.strip()

                    if actual_output == expected_output:
                        passed_cases += 1
                    else:
                        submission.result = 4  # 答案错误
                        break

                # 更新提交结果
                submission.time_used = max_time
                submission.memory_used = int(max_memory * 1024)  # 转换为KB
                submission.test_cases_passed = passed_cases
                
                # 如果所有测试用例都通过
                if passed_cases == submission.total_test_cases:
                    submission.result = 3  # 通过

            finally:
                # 清理容器
                try:
                    container.stop(timeout=10)  # Give container 10 seconds to stop gracefully
                except Exception as e:
                    print(f"Error stopping container: {e}")
                
                try:
                    container.remove(force=True)  # Force remove if needed
                except Exception as e:
                    print(f"Error removing container: {e}")

    except Exception as e:
        submission.result = 9  # 系统错误
        submission.remark = str(e)

    finally:
        submission.save()


def get_language_config(language):
    """
    获取语言配置
    """
    configs = {
        0: {  # C
            'image': 'gcc:latest',
            'file_name': 'main.c',
            'compile_cmd': ['gcc', '-o', '/app/program', '/app/main.c'],
            'run_cmd': ['/app/program'],
            'time_limit': 1000,  # 毫秒
            'memory_limit': 256  # MB
        },
        1: {  # C++
            'image': 'gcc:latest',
            'file_name': 'main.cpp',
            'compile_cmd': ['g++', '-o', '/app/program', '/app/main.cpp'],
            'run_cmd': ['/app/program'],
            'time_limit': 1000,
            'memory_limit': 256
        },
        3: {  # Python
            'image': 'python:3.9',
            'file_name': 'main.py',
            'run_cmd': ['python', '/app/main.py'],
            'time_limit': 2000,
            'memory_limit': 256
        }
    }
    return configs.get(language)